import torch
import torchvision.transforms as transforms
from PIL import Image
import torch.nn as nn

# 定义MLP模型
class MLP(nn.Module):
    def __init__(self, input_size, hidden1, hidden2, num_classes):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden1)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(hidden1, hidden2)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(hidden2, num_classes)

    def forward(self, x):
        out = self.relu1(self.fc1(x))
        out = self.relu2(self.fc2(out))
        out = self.fc3(out)
        return out

# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载模型
input_size = 224 * 224 * 3
hidden1 = 512
hidden2 = 256
num_classes = 3
model = MLP(input_size, hidden1, hidden2, num_classes).to(device)
model.load_state_dict(torch.load('mlp_best_model.pth'))  # 加载最佳模型权重
model.eval()

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 预测函数
def predict(image_path):
    # 加载并预处理图像
    image = Image.open(image_path).convert('RGB')
    image = transform(image)
    image = image.view(-1, 224*224*3).to(device)

    # 模型预测
    with torch.no_grad():
        outputs = model(image)
        _, predicted = torch.max(outputs, 1)
        predicted = predicted.cpu().numpy()

    # 类别名称
    class_names = ['最优级-1级', '中间级-2级', '差等级-3级']  # 类别名
    return class_names[predicted[0]]

# # 示例用法
# image_path = r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\test_img\test_605.tif'  # 替换为实际的图片路径
# predicted_class = predict(image_path)
# print("Predicted Class:", predicted_class)

image_path_list = [
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_1\1.tif',
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_2\1.tif',
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_3\1.tif'

]  # 替换为你的图片路径

for image_path in image_path_list:
    predicted_class = predict(image_path)
    print("Predicted Class:", predicted_class)